'''
Created on 2011-3-3

@author: summit
'''

import re
import numpy
import math
import array

def pgm_nextline(fin):
    while True:
        header  = fin.readline().strip()
        if header.startswith('#'):
            continue
        else:
            break
    return header
            
def pgm2numpy_p5(fin, debug = True):
    try:
        
        rows, cols = 0, 0
        raw_data = re.split('\W+', pgm_nextline(fin))
        rows = int(raw_data[1])
        cols = int(raw_data[0])
        raw_data = re.split('\W+', pgm_nextline(fin))
        bits = (int(raw_data[0])+1)/256
        
        assert (rows, cols) != (0,0)
        
        if debug:
            print 'Rows: %d, cols: %d' %(rows, cols)
        
        assert bits > 0
        
        if debug:
            print 'color bit is %d' % bits
        
        if bits<=1:
            binvalues = array.array('c') # unsigned char 
        elif bits<=2 and bits>1:
            binvalues = array.array('H') # unsigned short
        elif bits<=4 and bits>2:
            binvalues = array.array('f') # float
        else:
            print 'color bit is %d too big' % bits
        result = numpy.fromfile(fin, dtype=numpy.uint8)
        result = numpy.reshape(result, (rows, cols))
        
        return result
            
    finally:
        if fin != None:
            fin.close()
        fin = None
    
    return None
    
    
def pgm2numpy_p2(fin, debug = True):
    
    try:
        
        rows, cols = 0, 0
        
        raw_data = re.split('\W+', pgm_nextline(fin))
        
        rows = int(raw_data[1])
        cols = int(raw_data[0])
        raw_data = re.split('\W+', pgm_nextline(fin))
        bits = (int(raw_data[0])+1)/256
        
        assert (rows, cols) != (0,0)
        
        if debug:
            print 'Rows: %d, cols: %d' %(rows, cols)
        
        assert bits > 0
        
        if debug:
            print 'color bit is %d' % bits
        #
        # Initialise a 2D numpy array
        #
        if bits<=1:
            result = numpy.zeros((rows, cols), numpy.int8)
        elif bits<=2 and bits>1:
            result = numpy.zeros((rows, cols), numpy.int16)
        elif bits<=4 and bits>2:
            result = numpy.zeros((rows, cols), numpy.int32)
        else:
            print 'color bit is %d too big' % bits
        
        pxs = []
        raw_data = re.split('\W+', fin.read())
        #
        # Read to EOF.
        #
        for v in raw_data:
            if v == '':
                break
            
            pxs.append(int(v))
        
        if len(pxs) != rows*cols:
            if debug:
                print 'Insufficient image data:', len(pxs)
            return None
        
        for r in range(rows):
            for c in range(cols):
                #
                # Index into the numpy array and set the pixel value
                #
                result[r, c] = pxs[r*cols+c]
        return result
    
    finally:
        if fin != None:
            fin.close()
        fin = None
    
    return None

def pgm2numpy(filename):
    '''
    Read a PGM into a numpy array. 
    '''
    fin = None 
    debug = True
    r = None

    fin = open(filename, 'rb')
    header = pgm_nextline(fin)
 
    if header == 'P1':
        pass
    elif header == 'P2':
        if debug:
            print 'PBM is p2 format'
        r =  pgm2numpy_p2(fin, debug)
    elif header == 'P4':
        assert False, 'Raw PBM reading not implemented yet'
    elif header == 'P5':
        if debug:
            print 'PBM is p5 format'
        r =  pgm2numpy_p5(fin, debug)
    else:
        #
        # unexpected header.
        #
        if debug:
            print 'Bad mode: ', header
        return None
    return r    
        
if __name__ == '__main__':
    import matplotlib.pyplot as plt
    x = pgm2numpy('../data/mdb092.pgm')
    img = numpy.zeros((x.shape[0], x.shape[1], 3))
    img[:,:,0] = img[:,:,1] = img[:,:,2] = x    # matplotlib only should the RGB
    print img.shape
    plt.figure()
   
    plt.imshow(img)
   
    plt.show()